
/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2020-2021. All rights reserved.
 * Description: implement of DrawBbox
 */
#include "draw_bbox_kernels.h"

namespace  {
const char *DRAW_BBOX = "DrawBbox";
}

namespace aicpu  {
uint32_t DrawBboxCpuKernel::Compute(CpuKernelContext &ctx)
{
    //get input tensor
    Tensor *image_tensor = ctx.Input(0);
    if (image_tensor == nullptr) {
        return 1;
    }

    Tensor *image_info_tensor = ctx.Input(1);
    if (image_info_tensor == nullptr) {
        return 1;
    }

    Tensor *bbox_data_tensor = ctx.Input(2);
    if (bbox_data_tensor == nullptr) {
        return 1;
    }

    Tensor *bbox_num_tensor = ctx.Input(3);
    if (bbox_num_tensor == nullptr) {
        return 1;
    }

    Tensor *bbox_color_tensor = ctx.Input(4);
    if (bbox_color_tensor == nullptr) {
        return 1;
    }

    Tensor *bbox_thickness_tensor = ctx.Input(5);
    if (bbox_thickness_tensor == nullptr) {
        return 1;
    }

    //get image shape
    int32_t width = ((int32_t*)image_info_tensor->GetData())[0];
    int32_t height = ((int32_t*)image_info_tensor->GetData())[1];
    //get image data
    Mat img = Mat(height, width, CV_8UC3, (uchar*)image_tensor->GetData());
    int32_t* bbox_data = (int32_t*)bbox_data_tensor->GetData();
    int32_t bbox_num = ((int32_t*)bbox_num_tensor->GetData())[0];
    int32_t* bbox_color = (int32_t*)bbox_color_tensor->GetData();
    int32_t bbox_thickness = ((int32_t*)bbox_thickness_tensor->GetData())[0];

    for(int n = 0; n < bbox_num; n++){
        rectangle(img, Point2i(bbox_data[n * 4 + 0], bbox_data[n * 4 + 1]),
        Point2i(bbox_data[n * 4 + 2], bbox_data[n * 4 + 3]),
        Scalar(bbox_color[0], bbox_color[1], bbox_color[2]), bbox_thickness);
    }

    return 0;
}

REGISTER_CPU_KERNEL(DRAW_BBOX, DrawBboxCpuKernel);
} // namespace aicpu
